import random
from math import sqrt
from typing import List, Tuple, Union
import pygame

import gym
import numpy as np


class GridWorldEnv(gym.core.Env):
    W, H = 9, 7
    MOVE_UP, MOVE_DOWN, MOVE_LEFT, MOVE_RIGHT = range(4)
    TYPE_GRID, TYPE_WALL, TYPE_TRAP, TYPE_GOAL, TYPE_PLAYER = range(5)
    COLOR = [(255, 255, 255), (0, 0, 0), (255, 0, 0), (0, 255, 0), (255, 255, 0)]
    MAX_STEP = 500  # 最多走500步

    GRID_A = 80  # 正方形边长
    TRAP_REW = -600
    GOAL_REW = 100
    walls = [(1, 1), (2, 2), (2, 1), (6, 1), (7, 2), (7, 1), (3, 5), (4, 5), (5, 5)]
    reds = [(1, 2), (6, 2), (4, 3)]
    greens = [(4, 0)]
    action_space = gym.spaces.Discrete(4)
    observation_space = gym.spaces.Box(
        low=0, high=255, shape=(W, H), dtype=np.int
    )

    def draw_rect(self, color, x, y):
        pygame.draw.rect(self.screen, color, (x, y, self.GRID_A - 1, self.GRID_A - 1))

    def __init__(self):
        self.screen = pygame.display.set_mode((self.W * self.GRID_A, self.H * self.GRID_A))
        pygame.display.set_caption("GridWorld")
        pygame.init()

        self.grid_map, self.player = self.init()
        self.start_point = self.player
        self.total_step = 0

    def init(self) -> (np.ndarray, List[int]):
        grid_map = []
        for row in range(self.H):
            grid_map.append([self.TYPE_GRID] * self.W)
        for (col, row) in self.reds:
            grid_map[row][col] = self.TYPE_TRAP
        for (col, row) in self.greens:
            grid_map[row][col] = self.TYPE_GOAL
        for (col, row) in self.walls:
            grid_map[row][col] = self.TYPE_WALL

        player = [int(self.W / 2), self.H - 1]
        grid_map[player[1]][player[0]] = self.TYPE_PLAYER
        return np.array(grid_map), player

    def step(self, action: int):
        pygame.event.pump()  # 内部处理pygame事件处理程序
        reward, done = self.move(
            dx=1 if action == self.MOVE_RIGHT else (-1 if action == self.MOVE_LEFT else 0),
            dy=1 if action == self.MOVE_DOWN else (-1 if action == self.MOVE_UP else 0)
        )
        self.total_step += 1
        if self.total_step >= self.MAX_STEP:
            done = True
        return self.grid_map, reward, done, {}

    def move(self, dx, dy) -> (int, bool):
        x, y = self.player[0] + dx, self.player[1] + dy
        if not (0 <= x < self.W and 0 <= y < self.H):
            return -1, False
        grid_type = self.grid_map[y, x]
        if grid_type != self.TYPE_WALL:
            self.grid_map[self.player[1]][self.player[0]] = self.TYPE_GRID
            self.player = [x, y]
            self.grid_map[y][x] = self.TYPE_PLAYER
        if grid_type == self.TYPE_TRAP:
            return self.TRAP_REW, True
        if grid_type == self.TYPE_GOAL:
            return self.GOAL_REW, True
        return -1, False

    def reset(self):
        self.__init__()
        return self.grid_map

    def close(self):
        pygame.quit()

    def seed(self, seed=None):
        if seed is not None:
            random.seed(seed)
            return [seed]
        else:
            return []

    def render(self, mode='human'):
        for row in range(self.grid_map.shape[0]):
            for col in range(self.grid_map.shape[1]):
                self.draw_rect(self.COLOR[self.grid_map[row, col]], col * self.GRID_A, row * self.GRID_A)
        pygame.display.update()

    @staticmethod
    def get_action_meanings():
        return ["UP", "DOWN", "LEFT", "RIGHT"]

    @staticmethod
    def manhattan(p1: Union[Tuple[int, int], List[int]], p2: Union[Tuple[int, int], List[int]]):
        # 两点的曼哈顿距离
        return abs(p1[0] - p2[0]) + abs(p1[1] - p2[1])

    @staticmethod
    def distance(p1: Union[Tuple[int, int], List[int]], p2: Union[Tuple[int, int], List[int]]):
        # 两点的直线距离
        return sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)
